from brian import *
from brian.library.random_processes import *
from brian.library.synapses import  *
import MoNeuron as Moto
import RNeuron as Renshaw
import InNeuron as IN
from time import clock
from MyConstants import gSynE, gSynI
from matplotlib.backends.backend_pdf import PdfPages
from time import clock
from brian.tools.datamanager import *
from brian.tools.taskfarm import *


class recurrent_inhibition(object):
    def __init__(self):
        self.N = 50
        self.DUR = 3000 * ms
        
        
        self.MNFactory = Moto.MnNeuronFactory()
        self.MN = self.MNFactory.init(self.N, 'alpha', 1, False)
        self.RFactory = Renshaw.RNeuronFactory()
        self.Renshaw = self.RFactory.init(self.N, 'biexp', 1, True)
        
        
        self.spiketimes = []
        for i in xrange(0, self.N):
             self.spiketimes.append((i, 4 * ms + randn() * 2 * ms))
        
        self.AfferentInput = SpikeGeneratorGroup(self.N, self.spiketimes, period=1 / (16 * Hz))
        
        self.C_MN_RE = Connection(self.MN, self.Renshaw, 'ge', weight=5.0 / self.N * msiemens * cm ** -2, sparseness=0.8, delay=(1 * ms, 3 * ms))
        self.C_RE_MN = Connection(self.Renshaw, self.MN, 'gi', weight=0.05 / self.N * 0.05 * msiemens * cm ** -2, sparseness=0.5, delay=(1 * ms, 3 * ms))
        self.C_Aff_MN = Connection(self.AfferentInput, self.MN, 'ge', weight=8.0 / self.N * 0.05 * msiemens * cm ** -2, sparseness=1)
        
        self.spikeRE = SpikeMonitor(self.Renshaw)
        self.rateRE = PopulationRateMonitor(self.Renshaw,10*ms)
        self.spikeMN = SpikeMonitor(self.MN)
        
        self.traceMN = StateMonitor(self.MN,'vm',record=range(0,self.N))

        @network_operation(clock=EventClock(dt=200 * ms))
        def progress(clk):
            print clk._t
            return
        
        self.net = Network(self.MN, self.Renshaw, self.AfferentInput, self.C_MN_RE, self.C_RE_MN, self.C_Aff_MN, self.spikeRE, self.spikeMN,self.traceMN,self.rateRE, progress)
        clear(False,True)
        
    def __call__(self,run_id,w_aff_mn, w_mn_re, w_re_mn,report=0):
        self.C_MN_RE.__init__(self.MN, self.Renshaw, 'ge', weight=w_mn_re / self.N * msiemens * cm ** -2, sparseness=0.8, delay=(1 * ms, 3 * ms))
        self.C_RE_MN.__init__(self.Renshaw, self.MN, 'gi', weight=w_re_mn / self.N * 0.05 * msiemens * cm ** -2, sparseness=0.5, delay=(1 * ms, 3 * ms))
        self.C_Aff_MN.__init__(self.AfferentInput, self.MN, 'ge', weight=w_aff_mn / self.N * 0.05 * msiemens * cm ** -2, sparseness=1)
        
        self.net.reinit(states=True)
        self.MNFactory.reinit()
        self.RFactory.reinit()
        self.spikeRE.reinit()
        self.rateRE.reinit()
        self.spikeMN.reinit()
        self.traceMN.reinit()
        
       
   
        filename = 'recurrent_inhibition, id = ' + str(run_id) + ', w_aff_mn = ' + str(w_aff_mn) + ', w_mn_re = ' + str(w_mn_re) + ', w_re_mn = ' + str(w_re_mn) + '.pdf'
        
        start = clock()
        if report !=0:
            self.net.run(self.DUR,report=report)
        else:
            self.net.run(self.DUR)
        
        stop = clock()
        print 'Duration ' + str(stop - start)
        
        
        binsize = 1 / (2 * 16 * Hz)
        MNspikes = list()
        
        for x in self.spikeMN.spikes:
            MNspikes.append(x[1] / ms)
        
        
        pp = PdfPages(filename)
        
        fig1 = figure(1)
        subplot(211)
        bar(self.rateRE.times/ms,self.rateRE.rate,width=10,color='w')
        fig1.axes[0].set_xlim(0, self.DUR / ms)
        subplot(212)
        raster_plot(self.spikeMN)
        fig1.axes[1].set_xlim(0, self.DUR / ms)
        
        self.sum_mn = self.traceMN.__getitem__(0)
        for i in xrange(1, self.N):
            self.sum_mn = self.sum_mn + self.traceMN.__getitem__(i)
            
        self.sum_mn = self.sum_mn / self.N
        
        fig2 = figure(2)
        if self.spikeMN.spikes.__len__() != 0:
            subplot(211)
        plot(self.sum_mn)
        self.ret = 0
        if self.spikeMN.spikes.__len__() != 0:
            subplot(212)
            self.ret=hist(MNspikes, floor(self.DUR / binsize), range=[0, self.DUR / ms], histtype='step')
        
        self.consistency = 0
        self.strength = 0
        if self.ret !=0:
            responses = self.ret[0][::2][7:]
            uneven = zeros(int(floor(responses.__len__()/2.0)))
            even = zeros(int(floor(responses.__len__()/2.0)))
            for i in xrange(0,int(floor(responses.__len__()/2.0))):
                uneven[i]=(responses[i*2-1])
                even[i]=(responses[i*2])
                
            diff1 = uneven-even
            diff2 = uneven[1:uneven.__len__()]-even[0:even.__len__()-1]
            diff = concatenate((diff1,diff2),1)
            
            sdiff = sign(diff)
            cs = zeros(3)
            for i in xrange(0,diff.__len__()):
                if sdiff[i] == 1:
                    cs[0]+=1
                elif sdiff[i] == -1:
                    cs[1]+=1
                else:
                    cs[2]+=1
            
            bigger = cs[0] if cs[0]>cs[1] else cs[1]
            smaller = cs[0] if cs[0]<=cs[1] else cs[1]
            
            self.consistency = float(bigger)/float(smaller+bigger)
            self.consistency -= 0.5
            self.consistency *= 2
            
            self.strength = mean(diff)/mean(responses)    
            
            
        fig1.set_size_inches(40, 20)
        fig2.set_size_inches(40, 20)
        fig2.suptitle('consistency = ' + str(self.consistency) + ', strength = ' + str(self.strength), fontsize=12)
        
        
        
        pp.savefig(fig2)
        pp.savefig(fig1)
        
        pp.close()
        #show()
        close('all')
        
        
        return (run_id,(w_aff_mn, w_mn_re, w_re_mn,self.consistency,self.strength))

if __name__ == '__main__':
    #ri = recurrent_inhibition()
    #ri.__call__(1,8.0,5.0,0.05)
    dataman = DataManager('ri-3-5')
    
    run = 'vary2'
    if run == 'once':
        ri = recurrent_inhibition()
        ri.__call__(1,8.0,5.0,0.03)
    elif run == 'vary':
        b = linspace(3.0,7.0,11)
        c = linspace(0.025,0.075,11)
        
        cutoff = dataman.itemcount()-126
        input = []
        for i in xrange(0,11):
            for j in xrange(0,11):
                for k in xrange(0,5):
                    if ((i*11+j)*5+k)>cutoff:
                        input.append(((i*11+j)*5+k,8.0,b[i],c[j]))
        print input
        run_tasks(dataman, recurrent_inhibition, input)
    elif run == 'vary2':
        input = []
        for i in xrange(0,50):
            input.append((i,8.0,3.0,0.05))
        run_tasks(dataman, recurrent_inhibition, input)
    else:
        X, Y = zip(*dataman.values())
        le =(len(X))/5
        x = zeros(le)
        y = zeros(le)
        
        measures = dict()
        
        
        c=0
        for i in Y:
            
            k = (Y[c][1],Y[c][2])
            item=(0,0,0)
            if measures.has_key(k):
                item = measures.get(k)
            newitem = (item[0]+Y[c][3],item[1]+abs(Y[c][4]),item[2]+1)
            measures[k]=newitem
            c+=1
            
        length = len(measures.values())
        res = measures.values()
        keys = measures.keys()
        consistency = zeros(length)
        strength = zeros(length)
        fitness = zeros(length)
        final_result = []
        c = 0
        for o in res:
            consistency[c] = o[0]/float(o[2])
            strength[c] = o[1]/float(o[0])
            fitness[c] = consistency[c]*strength[c]
            final_result.append((keys[c], consistency[c],strength[c],fitness[c]))
            #print 'id = ' + str(keys[c]) + ', consistency = ' + str(consistency[c]) +', strength = ' + str(strength[c]) + ', fitness = ' + str(fitness[c])
            c+=1
       
        sres= sorted(final_result, key=lambda final_result: final_result[1],reverse=True)      
           
        for o in sres:
            print o      
                
#    plot(X, Y, '.')
#    xlabel('strength')
#    ylabel('count firing')
#    show()            
